import javassist.*;
import javassist.bytecode.*;
import javassist.bytecode.annotation.*;

import java.util.List;


public class NotNullInstrumenter {

    private Options options;
    private ConstPool constPool;

    public NotNullInstrumenter(Options options, ConstPool constPool) {
        this.options = options;
        this.constPool = constPool;
    }

    public void insrumentMethod(CtMethod ctMethod) throws CannotCompileException, NotFoundException {
        instrumentParamAnnotations(ctMethod);
        instrumentReturnAnnotations(ctMethod);


    }

    private void instrumentReturnAnnotations(CtMethod ctMethod) {
        MethodInfo methodInfo = ctMethod.getMethodInfo();

        AttributeInfo ai = ctMethod.getMethodInfo().getAttribute("RuntimeInvisibleAnnotations");
        AnnotationsAttribute annAttr = (AnnotationsAttribute)methodInfo.getAttribute("RuntimeInvisibleAnnotations");
        if (annAttr == null) {
            return;
        }


        CodeAttribute ca = methodInfo.getCodeAttribute();
        CodeIterator iter = ca.iterator();
        try {
            //ca.setMaxStack(ca.getMaxStack() + 1);

            while (iter.hasNext()) {
                int index = iter.next();
                int op = iter.byteAt(index);

                if (op == Opcode.ARETURN) {

                    CtClass ctIllegalArgumentException = options.getClassPool().get("java/lang/IllegalArgumentException");
                    String descriptorString = Descriptor.ofConstructor(new CtClass[0]);

                    
                    Bytecode b = new Bytecode(constPool);
                    b.add(Opcode.DUP);
                    b.add(Opcode.IFNONNULL);
                    b.add(0);
                    b.add(12);
                    b.add(Opcode.POP);
                    b.addNew(ctIllegalArgumentException);
                    b.add(Opcode.DUP);
                    b.addInvokespecial(ctIllegalArgumentException, "<init>", descriptorString);
                    b.add(Opcode.ATHROW);


                    iter.insert(index, b.get());
                }
            }

            ca.computeMaxStack();
            
        } catch (BadBytecode badBytecode) {
            badBytecode.printStackTrace();  //To change body of catch statement use File | Settings | File Templates.
        } catch (NotFoundException e) {
            e.printStackTrace();  //To change body of catch statement use File | Settings | File Templates.
        }


    }

    private void instrumentParamAnnotations(CtMethod ctMethod) throws NotFoundException, CannotCompileException {

        ParameterAnnotationsAttribute annAttr = null;
        for (AttributeInfo ai : (List<AttributeInfo>)ctMethod.getMethodInfo().getAttributes()) {
            if (ai instanceof ParameterAnnotationsAttribute) {
                annAttr = (ParameterAnnotationsAttribute)ai;
                break;
            }
        }
        if (annAttr == null) {
            return;
        }

        CtClass[] paramTypes = ctMethod.getParameterTypes();

        // instrument param annotations
        StringBuilder sb = new StringBuilder();
        Annotation[][] matrixAnnotations = annAttr.getAnnotations();
        for (int i = 0; i < matrixAnnotations.length; i++) {

            if (!Utils.isInstanceOf(paramTypes[i], "Object")) {
                continue;
            }


            boolean existMyAnnotation = false;
            for (Annotation oneAnnotation : matrixAnnotations[i]) {
                System.out.format("    %s\n", oneAnnotation.toString());
                if (oneAnnotation.getTypeName().equals("NotNull")) {
                    existMyAnnotation = true;
                    break;
                }
            }

            if (existMyAnnotation) {
                sb.append("if ($");
                sb.append(i + 1);
                sb.append(" == null) { throw new NullPointerException(); }");
            }
        }
        ctMethod.insertBefore(sb.toString());
    }
}
